#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time    : 2024/9/19 9:27
# @Author  : Tom_zc
# @FileName: prepare-env.py
# @Software: PyCharm
import os
import shutil
import sys
import re
import logging
import click
import json
import yaml
import hashlib
import subprocess
from pathlib import Path

CURR_DIR = Path(__file__).resolve()
ROOT_DIR = CURR_DIR.parent.parent.parent.parent
sys.path.append(ROOT_DIR.as_posix())

from scripts.utils.common import func_retry
from scripts.utils.obs_helper import ObsHelper
from scripts.utils.common import execute_cmd3

logging.basicConfig(level=logging.INFO,
                    format="%(asctime)s  %(levelname)s:%(message)s")

clone_dir = 'source'


def generate_sha256_file(filename):
    file_path = Path(filename)
    sha256_path = file_path.parent / (file_path.stem + '.sha256')

    try:
        if not file_path.is_file():
            logging.error(f"文件不存在: {filename}")
            sys.exit(1)
        sha256_hash = hashlib.sha256()
        with file_path.open('rb') as f:
            for chunk in iter(lambda: f.read(8192), b''):
                sha256_hash.update(chunk)

        checksum = sha256_hash.hexdigest()
        sha256_path.write_text(f"{checksum} *{file_path.name}\n")

        return str(sha256_path), sha256_path.name
    except Exception as e:
        logging.error(f"生成sha256文件失败: {e}")
        sys.exit(1)


class PrePareEnv:
    clone_cmd = "git clone -b {} https://gitcode.com/openUBMC/{}.git"

    def prepare_source(self, repo, branch):
        if os.path.exists(repo):
            logging.info("WARNING: %s already exist, delete it." % repo)
            shutil.rmtree(repo)
        ret, out, err = execute_cmd3(
            self.clone_cmd.format(branch, repo))
        if ret != 0:
            logging.info("Error: Failed to git clone {}, err:{}, out:{}".format(repo, err, out))
            return 1
        return 0


class PipeLineLogHelper:

    def __init__(self, ak, sk, obs_endpoint, obs_bucket_name):
        self.obs_helper = ObsHelper(ak, sk, obs_endpoint)
        self.obs_bucket_name = obs_bucket_name

    def _get_bucket_key(self, prefix_path, filename):
        return "{}{}".format(prefix_path, filename)

    # prefix_path, zip_file_name, zip_file_name
    @func_retry()
    def upload_obs(self, prefix_path, path, filename):
        bucket_key = self._get_bucket_key(prefix_path, filename)
        logging.info("upload to obs: {}".format(bucket_key))
        self.obs_helper.upload_file(self.obs_bucket_name, bucket_key, path)


def get_manifest_components(manifest_dir: str):
    subsys_stable_dir = os.path.join(manifest_dir, "build", "subsys", "stable")
    sys_ymls = os.listdir(subsys_stable_dir)
    components = {}
    for sys_yml in sys_ymls:
        yml_path = os.path.join(subsys_stable_dir, sys_yml)
        with open(yml_path, "r") as fp:
            data = yaml.safe_load(fp)
        for pkg in data.get("dependencies", []):
            conan = pkg.get("conan")
            comp_name, comp_version, *_ = re.split("@|/", conan)
            components[comp_name] = comp_version
    return components


@click.command()
@click.option("--ak")
@click.option("--sk")
@click.option("--obs_endpoint")
@click.option("--bucket_name")
@click.option("--prefix_path")
@click.option("--zip_file_name")
def main(ak, sk, obs_endpoint, bucket_name, prefix_path, zip_file_name):
    if not all([ak, sk, obs_endpoint, bucket_name, prefix_path, zip_file_name]):
        logging.info("Lack of params, Please check.")
        sys.exit(-1)

    current_directory = os.path.dirname(os.path.abspath(__file__))
    json_file_path = os.path.join(current_directory, 'repo.json')
    with open(json_file_path, 'r', encoding='utf-8') as file:
        data = json.load(file)
    original_dir = os.getcwd()
    os.chdir(clone_dir)

    components = get_manifest_components(os.path.join(os.getcwd(), 'manifest'))
    for entry in data:
        repo = entry.get('repo')
        component = entry.get('component', repo)
        branch = components.get(component, entry.get('branch'))

        if repo and branch:
            result = PrePareEnv().prepare_source(repo, branch)
            if result != 0:
                sys.exit(-1)

    os.chdir(original_dir)
    command = ['zip', '-ry', zip_file_name, clone_dir]
    subprocess.run(command, cwd=original_dir, check=True)
    logging.info("=" * 20 + "end to prepare env" + "=" * 20)
    pipeline_log_helper = PipeLineLogHelper(ak, sk, obs_endpoint, bucket_name)
    pipeline_log_helper.upload_obs(prefix_path, zip_file_name, zip_file_name)
    sha256, sha256_filename = generate_sha256_file(zip_file_name)
    pipeline_log_helper.upload_obs(prefix_path, sha256, sha256_filename)
    sys.exit(0)


if __name__ == '__main__':
    main()
